from socket import *
import os
import signal
import sys
import time
# 文件库路径 
FILE_PATH = "/home/" 

class TftpServer(object):
    def __init__(self, connfd):
        self.connfd = connfd

    def do_list(self):
        # 获取列表
        file_list = os.listdir(FILE_PATH)
        # 如果对应的路径内没有文件,返回Empty
        if not file_list:
            self.connfd.send('Empty'.encode())
            return
            # 路径存在文件,向客户端发送OK
        else:
            self.connfd.send(b'OK')
            time.sleep(0.1)

        files = ""
        for file in file_list:
            # 排除以'.'开头的隐藏文件
            if file[0] != '.' and \
                    os.path.isfile(FILE_PATH + file):
                files = files + file + '#'
                # 返回文件列表
        self.connfd.send(files.encode())

        # 下载文件功能

    def do_get(self, filename):
        try:
            fd = open(FILE_PATH + filename, 'rb')
        except:
            self.connfd.send("File doesn't exist".encode())
            return
            # 如果能正常打开,发送OK
        self.connfd.send(b"OK")
        time.sleep(0.1)
        # 开始发送文件
        try:
            for line in fd:
                self.connfd.send(line)
            fd.close()
        except Exception as e:
            print(e)
        time.sleep(0.1)
        self.connfd.send(b'##')
        print("File send over")

        # 开始上传文件

    def do_put(self, filename):
        try:
            fd = open(FILE_PATH + filename, 'w')
        except:
            self.connfd.send("Some error")
            # 如果能正常打开文件,则发送OK
        self.connfd.send(b'OK')
        # 开始发送
        while True:
            # data为文件内容
            data = self.connfd.recv(1024).decode()
            if data == "##":
                break
            fd.write(data)
        fd.close()
        print("上传完毕")

def main():
        # 创建套接字/地址/端口
        HOST = '0.0.0.0'
        PORT = 8888
        ADDR = (HOST, PORT)

        sockfd = socket()
        # 设置端口可重用
        sockfd.setsockopt(SOL_SOCKET, SO_REUSEADDR, 1)
        # 绑定地址
        sockfd.bind(ADDR)
        # 设置监听队列大小
        sockfd.listen(5)

        signal.signal(signal.SIGCHLD, signal.SIG_IGN)
        print("Listen to port 8888....")

        while True:
            try:
                connfd, addr = sockfd.accept()
            except KeyboardInterrupt:
                sockfd.close()
                sys.exit("Server exit")
            except Exception as e:
                print(e)
                continue
            print("Client login:", addr)
            # 创建父子进程
            pid = os.fork()
            if pid < 0:
                print("Process creation failed")
                continue
            elif pid == 0:
                # 子进程负责请求接收和发送,所以节省资源,关闭连接套字
                sockfd.close()
                tftp = TftpServer(connfd)
                # 接收客户端请求
                while True:
                    data = connfd.recv(1024).decode()
                    if not data:
                        continue
                    # 调用do_list方法获取文件列表
                    elif data[0] == 'L':
                        tftp.do_list()
                    # data ==> G filename
                    # 文件名以G开头,以空格为间隔发送过来
                    elif data[0] == 'G':
                        filename = data.split(' ')[-1]
                        tftp.do_get(filename)
                    elif data[0] == 'P':
                        filename = data.split(' ')[-1]
                        tftp.do_put(filename)
                    elif data[0] == 'Q':
                        print("客户端退出")
                        sys.exit(0)

            else:
                connfd.close()
                continue


if __name__ == "__main__":
 main()
